import re
from ltp import LTP
from graphviz import Digraph
import numpy as np
from queue import Queue

dep_dict = {
    'SBV':1,
    'VOB':2,
    "IOB":3,
    "FOB":4,
    "DBL":5,
    "ATT":6,
    "ADV":7,
    "CMP":8,
    "COO":9,
    'POB':10,
    "LAD":11,
    "RAD":12,
    "IS":13,
    "HED":14,
    "WP":15
}
def is_digital_in_str(str):
    return bool(re.search(r'\d', str))

def get_num_by_list(str):
    sum = 0
    for data in str:
        # print(data)
        sum = sum*10 + int(data)
    return sum

def get_core_verb_index_by_dep(deps):
    index = -1
    for dep in deps:
        # print('-->',dep)
        if dep[2] == "HED":
            index = dep[0]
    return index

def get_deps_view(words,deps,fig_name='测试图片',Debug=True):
    rely_id = [t[1] for t in deps] # 提取父节点ID
    relation = [ t[2] for t in deps] # 提取依赖关系
    heads = ['Root' if id==0 else words[id-1] for id in rely_id]


    g = Digraph(name=fig_name,format='png')

    g.node(name = "Root")
    for index,word in enumerate(words):
        temp = word
        g.node(name=temp,fontname = 'Fangsong')
    
    for i in range(len(words)):
        if relation[i] not in ["HED"]:
            g.edge(words[i],heads[i],label =relation[i])
        else:
            if heads[i] == "Root":
                g.edge(words[i],'Root',label =relation[i])
            else:
                g.edge(heads[i],'Root',label =relation[i])
    g.render(directory='.',view=Debug)

'''获取动词的主谓关系与动宾关系'''
def get_verb_sbv_vob(words,deps,debug=False):
    '''
    数据格式：
    (9, 0, 'HED'),0 - > 9 ，第二个为起点，第一个为终点 
    (8, 9, 'SBV'), 9 - > 8
    (11, 9, 'VOB'), 9 - > 11
    (7, 8, 'ATT'), 8 -> 7
    (10, 11, 'ATT'), 11 -> 10
    '''
    if debug:
        print(words[0])
        print(deps[0])
    # 先找核心动词
    for dep in deps[0]:
        if dep[2] == "HED":
            core_verb = dep
    print("核心动词为：",words[0][core_verb[0]-1])
    sbv_list = []
    vob_list = []
    # 再找主谓关系和动宾关系
    for dep in deps[0]:
        if dep[1] == core_verb[0] and dep[2] == "SBV":
            if debug: print("找到主谓关系：",dep)
            sbv_list.append(dep)
        if dep[1] == core_verb[0] and dep[2] == "VOB":
            if debug: print("找到动宾关系",dep)
            vob_list.append(dep)
    need_list = sbv_list + vob_list
    print(need_list)
    # 寻找定中关系
    for need_dep in need_list:
        for dep in deps[0]:
            if need_dep[0] == dep[1] and dep[2] == "ATT":
                x = dep[0] - 1
                y = dep[1] - 1
                print(dep,words[0][x],words[0][y])

    return sbv_list, vob_list

class Graph:
    def __init__(self,segs,deps,poses,debug=False):
        self.mat = None # 邻接矩阵，当中的节点下标减一
        self.n = len(segs)   # 节点数量
        self.arc_num = len(deps) -1 # 边的数量
        self.segs = segs
        self.deps = deps
        self.poses = poses
        self.visited = [False] * self.n # 标识当前节点是否被访问
        self.debug = debug

        self.att_completion = [] # 定语补全
        

        self.core_verb = self.get_core_verb()
        self.core_word = self.segs[self.core_verb-1]
        # print("查看核心动词：",self.core_verb)
        # print("节点数量：",self.n)
        # print("图初始化完成\n")
    
    def get_core_verb(self):
        temp = -1
        for dep in self.deps:
            if dep[2] == "HED":
                temp = dep[0]
        return temp
    def get_verb_sbv_vob(self):
        '''寻找主谓关系'''
        sbv = 0
        vob = 0
        for dep in self.deps:
            if dep[1] == self.core_verb and dep[2]=="SBV":
                sbv = dep[0]
            elif dep[1] == self.core_verb and dep[2] == "VOB":
                vob = dep[0]
        if self.debug:
            print("主谓关系：",sbv)
            print("动宾关系：",vob)
        return sbv,vob

    def createMat(self):
        '''根据依赖关系创建邻接矩阵'''
        # print(self.deps)
        tmp_mat = np.zeros((self.n,self.n))
        # print("邻接矩阵形状：",tmp_mat.shape)
        
        # 将依存关系存入图中
        for dep in self.deps:
            end,start,type = dep
            if start == 0 : continue # 跳过从0出发的虚拟节点
            tmp_mat[start-1][end-1] = dep_dict[type]
            if self.debug:
                print("终点 {} 起点 {} 类型 {} {}".format(end,start,type,dep_dict[type]))
            
        self.mat = tmp_mat
        # print("邻接矩阵创建完成\n")
        # print(np.count_nonzero(self.mat),self.arc_num)
        
    def dfs(self,verId,edge):
        '''从指定节点开始遍历'''
        self.visited[verId-1] = True
        self.att_completion.append(verId)
        if self.debug:
            print("当前节点：",verId,self.segs[verId-1])
        for w in range(self.n): # 指定边的类型
            if self.mat[verId-1][w] != 0 and self.mat[verId-1][w] in edge and not self.visited[w]:
                self.dfs(w+1,edge)
    
    def att_completion_by_dfs(self):
        '''完成实体的定语补全，返回第一个实体的补全结果'''
        sbv, _ = self.get_verb_sbv_vob() # 获得主语，宾语，同核心动词发生关系的
        
        self.att_completion.clear() # 清除数据
        temp_list_1 = [dep_dict["ATT"],dep_dict["COO"]]
        # print(temp_list_1)
        self.dfs(sbv,temp_list_1)
        sbv_list = self.att_completion # 保存数据
        sbv_list_char = [self.segs[t-1] for t in sbv_list]
        if self.debug:
            pass
            # print(sbv_list)
            # print('==>',sbv_list_char)
        sbv_entity = ''.join(reversed(sbv_list_char))
        return sbv_entity

    def vob_complete_by_bfs(self,coo_verb=0,search_obj=None):
        '''完成实体的宾语补全，返回第二个实体的补全结果'''
        search_list = [dep_dict["VOB"]]
        if coo_verb == 0:
            _,vob = self.get_verb_sbv_vob()
        else:
            vob = coo_verb # 从指定位置开始查找
            search_list.append(search_obj)
        # print("搜索列表",search_list)
        self.visited = [False] * self.n # 初始化
        # print("==============>",vob)
        Q = Queue()
        Q.put(vob)
        self.visited[vob-1] = True  # 标记访问
        vob_list = [vob-1]
        while not Q.empty():
            u = Q.get()
            for w in range(self.n):
                if self.mat[u-1][w] !=0 and not self.visited[w]  \
                    and self.mat[u-1][w] in search_list :
                    # print("-->",w,self.segs[w])
                    self.visited[w] = True
                    Q.put(w+1)
                    vob_list.append(w)
        if coo_verb != 0:
            vob_list = vob_list[1:]
        vob_list = [w for w in vob_list if self.poses[w] != "wp"]
        char_vob_list = [self.segs[w] for w in vob_list ]
        # print("-*>",vob_list) # 这里是词表的下标
        # print("-*>",char_vob_list) # 这里是词表的下标
        vob_completed = []
        # 下面对宾语列表进行att补全
        for t_vob in vob_list:
            self.att_completion.clear()
            temp_list = [dep_dict["ATT"]]
            self.dfs(t_vob+1,temp_list)
            # print('-->',self.att_completion)
            # print('-->',sorted(self.att_completion))
            vob_completed += sorted(self.att_completion) # 按从小到大排序
        char_vob_completed = [self.segs[t-1] for t in vob_completed]
        vob_entity = ''.join(char_vob_completed)
        # print(''.join(char_vob_completed))
        return vob_entity
    def get_coo_verb_by_core(self):
        temp = 0
        count = 0
        for dep in self.deps:
            if dep[2] == "COO" and dep[1] == self.core_verb:
                temp = dep[0]
                count += 1

        if count > 1: print("出现了多次与核心动词并列的动词")

        return temp

    def check_coo_verb_svb_vob(self):
        '''检查是否有词语与并列动词存在svb或vob结构'''
        coo_verb = self.get_coo_verb_by_core()
        coo_svb = 0
        coo_vob = 0
        for dep in self.deps:
            if dep[1] != coo_verb: continue
            if dep[2] == "SVB": 
                coo_svb = dep[0]
            elif dep[2] == "VOB":
                coo_vob = dep[0]
        
        print("svb vob 存在情况分别为：",coo_svb,coo_vob)
        return coo_svb,coo_vob

    def check_coo_for_svb_vob(self):
        '''
        检查是否有并列的主语或者并列的宾语
            如果有检测到并列的主语，则客体相同
            如果有检测到并列的宾语，则主体相同
        '''
        svb,vob = self.get_verb_sbv_vob()
        coo_svb = []
        coo_vob = []

        for dep in self.deps:
            if dep[1] == svb and dep[2] == "COO":
                coo_svb.append(dep[0])
            elif dep[1] == vob and dep[2] == "COO":
                coo_vob.append(dep[0])

        word_coo_svb = [self.segs[t-1] for t in coo_svb]
        word_coo_vob = [self.segs[t-1] for t in coo_vob]
        if self.debug:
            print("svb并列数量：{}, vob并列数据：{}".format(
                len(coo_svb),len(coo_vob)))
            print(word_coo_svb,word_coo_vob)
        return word_coo_svb,word_coo_vob

    def check_cmp_core_verb(self):
        '''检查主谓介宾关系，返回'''
        temp_tmp = 0
        for dep in self.deps:
            if dep[1] == self.core_verb and dep[2] == "CMP":
                temp_tmp = dep[0]

        if temp_tmp == 0: 
            print("没有找到动补结构，返回")
            return
        
        entity_list = []
        f_entity = self.att_completion_by_dfs() # 获得第一个实体
        # print("\n找到动补结构",self.segs[temp_tmp-1],temp_tmp)
        realton = self.core_word + self.segs[temp_tmp-1]
        # print("实体关系为：",realton)
        extra_entity = self.vob_complete_by_bfs(coo_verb=temp_tmp,search_obj=dep_dict["POB"])
        # print("找到的实体为：",extra_entity)
        # print("知识三元组为：",f_entity,realton,extra_entity)
        entity_list.append((f_entity,realton,extra_entity)) 
        return entity_list

    def parall_situation1(self):
        '''用于出现与核心动词并列的动词的情况'''
        coo_verb = self.get_coo_verb_by_core() # 获取与核心动词并列的动词
        if coo_verb == 0:
            print("第一种并列情况未发生，故不处理")
            return
        # print("coo_verb",coo_verb)
        # 取出与核心动词相关的实体
        extra_entity = self.vob_complete_by_bfs(coo_verb=coo_verb)
        f_entity = self.att_completion_by_dfs() # 得到的是字符串
        e_entity = self.vob_complete_by_bfs() # 得到的是字符串
        
        # print("f_entity",f_entity)
        # print("e_entity",e_entity)
        # print("extra_entity",extra_entity)
        entity = (f_entity,self.core_word,e_entity)
        all_entities_list = []
        all_entities_list.append(entity)
        all_entities_list.append((f_entity,self.core_word,extra_entity))
        # 添加额外的vob实体
        # for e_entity in extra_entity:
            # all_entities_list.append((entity[0],entity[1],e_entity))
        
        # print("-------> 这里是并列情况1")
        # print(all_entities_list)
        return all_entities_list

    def parall_situation2(self):
        # print("这里是并行情况2")
        f_entity = self.att_completion_by_dfs() # 获得头实体
        e_entity = self.vob_complete_by_bfs() # 获得尾实体
        if e_entity == "":
            print("这里是并行情况2，未找到尾实体，直接return")
            return
        entity = (f_entity,self.core_word,e_entity)
        print(entity)
        print("------------------------------")
        # 检查是否有并列的宾语或主语
        word_coo_svb,word_coo_vob = self.check_coo_for_svb_vob()

        # print(word_coo_svb)
        # print("======================>")
        # print(word_coo_vob)
        all_entities_list = []
        all_entities_list.append(entity)
        #---存在并列主语，则共享宾语
        for w_svb in word_coo_svb:
            all_entities_list.append((w_svb,entity[1],entity[2]))
        #---存在并列宾语，则共享主语
        for w_vob in word_coo_vob:
            all_entities_list.append((entity[0],entity[1],w_vob))
        # print("=========================> 这里是并列情况2")
        # print(all_entities_list)
        return all_entities_list

    def general_processing(self):
        # 1. 先检查存不存在主谓介宾结构
        first_entity_list = self.check_cmp_core_verb()

        # 2. 进行与核心动词并列的动词检查
        second_entity_list = self.parall_situation1()

        # 3. 检查是否有并列的主语或者宾语
        third_entity_list = self.parall_situation2()

        print("得到的实体1")
        print(first_entity_list)
        print("得到的实体2")
        print(second_entity_list)
        print("得到的实体3")
        print(third_entity_list)

def test1():
    '''
    测试定语补全与宾语补全
    '''
    sent = "变电站内主变、开关等重要设备发生严重故障。"
    sent2 = "调控机构负责编制电力统计。"
    ltp = LTP()

    seg, hidden = ltp.seg([sent])
    dep = ltp.dep(hidden)

    print(len(seg[0]))
    print(seg[0])
    g = Graph(seg[0],dep[0],debug=True)
    g.createMat()
 
    f_entity = g.att_completion_by_dfs()
    e_entity = g.vob_complete_by_bfs()
    print(f_entity,seg[0][g.core_verb-1],e_entity)





'''
coo 并列关系
'''
def test2():
    '''测试并列关系'''
    pass
if __name__ == "__main__":
    # print(get_num_by_list(['1','2','3']))
    sents = [
        '变电站内主变、开关等重要设备发生严重故障。',
        "调控机构负责编制电力统计。",
        '相关人员应熟悉厂站设备， 熟悉启动调试方案。',
        "发电厂应按照调度指令进行调频、调峰、调压。",
        "电厂调频能力应保持在系统频率在50±0.2Hz以内。"
    ]
    sent3 = "相关人员应熟悉厂站设备， 熟悉启动调试方案。"
    sent4 = "发电厂应按照调度指令进行调频、调峰、调压。"
    # sent5 测试主谓介宾结构
    sent5 = '电厂调频能力应保持在系统频率在50±0.2Hz以内。'

    # ltp = LTP()
    # # ltp.init_dict(path="user_dict.txt",max_window=4)
    # ltp.add_words(words=["发电厂","厂站"],max_window=4)
    # seg,hidden = ltp.seg([sent5])
    # pos = ltp.pos(hidden)
    # dep = ltp.dep(hidden)
    # seg_cont = seg[0]
    # dep_count = dep[0]
    # pos_cont = pos[0]
    # get_deps_view(seg_cont,dep_count,fig_name="test2")
    # print(seg_cont)
    # print(dep_count)
    # g = Graph(seg_cont,dep_count,pos_cont,debug=False)
    # g.createMat() # 首先需要创建邻接矩阵
    # # g.parall_situation1() # 并列结构的第1种情况
    # # g.parall_situation2() # 并列结构的第2种情况

    # # 开始测试主谓介宾结构
    # # g.check_cmp_core_verb()

    # # 测试三种情况
    # g.general_processing()
    ltp = LTP()
    ltp.add_words(words=["发电厂","厂站"],max_window=4)
    for index,sent in enumerate(sents):
        print("index  ------------------------------------------>")
        seg,hidden = ltp.seg([sent])
        pos = ltp.pos(hidden)
        dep = ltp.dep(hidden)
        seg_cont = seg[0]
        dep_cont = dep[0]
        pos_cont = pos[0]
        get_deps_view(seg_cont,dep_cont,fig_name=str(index),Debug=False)
        g = Graph(seg_cont,dep_cont,pos_cont,debug=False)
        g.createMat() # 创建邻接矩阵

        g.general_processing() # 进行通用处理
        print()
        



    

